#! /Pub/Apps/bin/python
import anndata
import typing

def cell_communication_by_c2c(adata: typing.Union[str, anndata.AnnData],
                              lr_path="/Pub/Users/wangyk/project/Poroject/F241121002_20250210/out/1.sc/5.cell_commutication/cell2cell/Human-2020-Jin-LR-pairs.csv",
                              od='./',
                              group_meta_path = '/Pub/Users/wangyk/project/Poroject/F241121002_20250210/out/1.sc/5.cell_commutication/cell2cell/group_meta_ch.csv',
                              cell_type_col='major_celltype',
                              focus_cell_type: list =['Neutrophil']):
    """
    使用cell2cell module 进行细胞通讯分析的函数。

    Args:
        adata (typing.Union[str, anndata.AnnData]): 输入的单细胞数据，可以是.h5ad文件的路径或anndata对象。
        lr_path (str): 配体-受体对数据的CSV文件路径，默认为指定路径。
        od (str): 输出目录，默认为当前目录。
        group_meta_path (str): 分组元数据的CSV文件路径，默认为指定路径。
        cell_type_col (str): 细胞类型列的名称，默认为'major_celltype'。
        focus_cell_type (str): 关注的细胞类型，默认为'Neutrophil'。

    Returns:
    无，将分析结果保存为PDF文件。
    """
    import scanpy as sc
    import os
    import cell2cell as c2c
    import pandas as pd
    import matplotlib.pyplot as plt
    import logging
    # 配置日志记录
    logging.root.handlers = []    
    logging.basicConfig(level=logging.INFO, format='\n[%(asctime)s] - %(levelname)s - %(message)s')
    
    os.makedirs(od, exist_ok=True)
    
    # 读取分组元数据
    group_meta = pd.read_csv(group_meta_path)
    
    # 获取分组颜色
    colors = c2c.plotting.get_colors_from_labels(labels=group_meta['Group'].unique().tolist(),
                                             cmap='Paired')

    # 根据adata类型读取数据
    if isinstance(adata, str):
        rnaseq = sc.read_h5ad(adata_path)
    elif isinstance(adata, anndata.AnnData):
        rnaseq = adata
    else:
        raise ValueError(
            "adata must be a str(path to your .h5ad file) or anndata.AnnData object")

    # 读取配体-受体对数据
    lr_pairs = pd.read_csv(lr_path)
    # 将数据转换为字符串类型
    lr_pairs = lr_pairs.astype(str)

    # 复制观测数据
    meta = rnaseq.obs.copy()
    # 细胞类型列名
    cell_type_col = cell_type_col

    logging.info("form c2c.analysis.SingleCellInteractions get interactions")
    # 创建单细胞相互作用对象
    interactions = c2c.analysis.SingleCellInteractions(
        rnaseq_data=rnaseq.to_df().T,
        ppi_data=lr_pairs,
        metadata=meta,
        interaction_columns=(
            'ligand_symbol', 'receptor_symbol'),
        communication_score='expression_thresholding',
        expression_threshold=0.1,  # values after aggregation
        cci_score='bray_curtis',
        cci_type='undirected',
        aggregation_method='nn_cell_fraction',
        barcode_col='index',
        celltype_col=cell_type_col,
        complex_sep='&',
        verbose=True)
    # 计算成对的通讯得分
    interactions.compute_pairwise_communication_scores()
    # 计算成对的细胞间相互作用得分
    interactions.compute_pairwise_cci_scores()

    # heatmap plot
    logging.info("form c2c.plotting.clustermap_ccc get interaction_clustermap")
    # 绘制聚类热图
    interaction_clustermap = c2c.plotting.clustermap_ccc(
        interactions,
        metric='jaccard',
        method='complete',
        metadata=group_meta,
        sample_col='Celltype',
        group_col='Group',
        colors=colors,
        row_fontsize=14,
        title='Active ligand-receptor pairs for interacting cells',
        filename=None,
        cell_labels=(
            'SENDER-CELLS', 'RECEIVER-CELLS'),
        **{'figsize': (10, 9),
           }
    )
    # Add a legend to know the groups of the sender and receiver cells:
    # 生成图例
    l1 = c2c.plotting.generate_legend(
        color_dict=colors,
        loc='center left',
        # Indicated where to include it
        bbox_to_anchor=(20, -2),
        ncol=1, fancybox=True,
        shadow=True,
        title='Groups',
        fontsize=14,
    )
    # 保存热图
    plt.savefig(f"{od}/0.all_cell_heatmap.pdf", bbox_inches='tight', dpi=300)

    # dot plot
    logging.info(
        "form interactions.permute_cell_labels get cci_pvals, ccc_pvals")
    # 计算细胞间相互作用的p值
    cci_pvals = interactions.permute_cell_labels(
        evaluation='interactions',
        permutations=10,
        fdr_correction=False,
        verbose=True)

    # 计算通讯的p值
    ccc_pvals = interactions.permute_cell_labels(
        evaluation='communication',
        permutations=10,
        fdr_correction=False,
        verbose=True)

    # 获取发送者和接收者细胞类型列表
    sender_cells = receiver_cells = meta[cell_type_col].unique().tolist()

    logging.info("form c2c.plotting.dot_plot get fig")
    # 绘制点图
    fig = c2c.plotting.dot_plot(
        interactions,
        evaluation='interactions',
        significance=0.2,
        figsize=(4, 7),
        cmap='PuOr',
        senders=sender_cells,
        receivers=receiver_cells,
        tick_size=9
    )
    # 保存点图
    plt.savefig(f"{od}/1.all_cell_interactions.pdf",
                bbox_inches='tight', dpi=300)

    # 获取所有细胞类型列表
    all_cell_type = meta[cell_type_col].unique().tolist()
    # 关注的细胞类型
    focus_cell_type = focus_cell_type
    # 绘制点图，以关注的细胞类型为接收者
    fig = c2c.plotting.dot_plot(
        interactions,
        evaluation='communication',
        significance=0.2,
        figsize=(5, 7.5),
        cmap='PuOr',
        senders=all_cell_type,
        receivers=focus_cell_type,
        tick_size=12
    )
    # 保存点图
    plt.savefig(f"{od}/1.{focus_cell_type}_as_receivers_.all_cell_communication.pdf",
                bbox_inches='tight', dpi=300)

    # 绘制点图，以关注的细胞类型为发送者
    fig = c2c.plotting.dot_plot(
        interactions,
        evaluation='communication',
        significance=0.2,
        figsize=(5, 7.5),
        cmap='PuOr',
        senders=focus_cell_type,
        receivers=all_cell_type,
        tick_size=12
    )
    # 保存点图
    plt.savefig(f"{od}/1.{focus_cell_type}_as_sender_.all_cell_communication.pdf",
                bbox_inches='tight', dpi=300)

    # 处理配体-受体对的p值数据
    df_of_cc_ligands_receptors = ccc_pvals.index.to_series().astype(
        str).str.strip('()').str.split(', ', expand=True)
    df_of_cc_ligands_receptors.columns = ['ligand', 'receptor']

    # 获取出现频率最高的前20个配体
    ligands = df_of_cc_ligands_receptors.iloc[:, 0].astype(
        str).str.strip("'").value_counts().head(20).index.tolist()
    # 获取出现频率最高的前20个受体
    receptors = df_of_cc_ligands_receptors.iloc[:, 1].astype(
        str).str.strip("'").value_counts().head(20).index.tolist()

    # 取交集，筛选出有效的配体和受体
    ligands = set(ligands) & set(lr_pairs.loc[:, 'ligand'].unique())
    receptors = set(receptors) & set(lr_pairs.loc[:, 'receptor'].unique())

    logging.info("form c2c.plotting.circos_plot get circos_plot")
    # 绘制环形图，以关注的细胞类型为发送者

    try:
        c2c.plotting.circos_plot(
            interaction_space=interactions,
            sender_cells= focus_cell_type,
            receiver_cells= all_cell_type,
            ligands=list(ligands),
            receptors=list(receptors),
            excluded_score=0,
            fontsize=16,
            ligand_label_color='black',
            receptor_label_color='brown',
        )
        # 保存环形图
        plt.savefig(f"{od}/2.circos_plot.focus_cell_type_as_sender.pdf",
                    bbox_inches='tight', dpi=300)

        # 绘制环形图，以关注的细胞类型为接收者
        c2c.plotting.circos_plot(
            interaction_space=interactions,
            sender_cells=all_cell_type,
            receiver_cells=focus_cell_type,
            ligands=list(ligands),
            receptors=list(receptors),
            excluded_score=0,
            fontsize=16,
            ligand_label_color='black',
            receptor_label_color='brown',
        )
        # 保存环形图
        plt.savefig(f"{od}/2.circos_plot.focus_cell_type_as_receiver.pdf",
                    bbox_inches='tight', dpi=300)

    except Exception as e:
        print(f"发生了一个错误:\n {e}")